/*
 * SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef POSTPROCESS_H_
#define POSTPROCESS_H_

#include <vector>
#include <iostream>
#include <string>
#include <memory>
#include "common/check.hpp"
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime_api.h>
// #include "iou3d_nms.h"
// #include "kernel.h"

/*
box_encodings: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading or *[cos, sin], ...]
anchors: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...]
*/
// struct Bndbox {
//     float x;
//     float y;
//     float z;
//     float w;
//     float l;
//     float h;
//     float rt;
//     int id;
//     float score;
//     Bndbox(){};
//     Bndbox(float x_, float y_, float z_, float w_, float l_, float h_, float rt_, int id_, float score_)
//         : x(x_), y(y_), z(z_), w(w_), l(l_), h(h_), rt(rt_), id(id_), score(score_) {}
// };

struct Bndbox {
    float x;
    float y;
    float z;
    float l;
    float w;
    float h;
    float rt;
    int id;
    float score;
    Bndbox(){};
    Bndbox(float x_, float y_, float z_, float l_, float w_, float h_, float rt_, int id_, float score_)
        : x(x_), y(y_), z(z_), l(l_), w(w_), h(h_), rt(rt_), id(id_), score(score_) {}
};
// struct Bndbox {
//     float x;
//     float y;
//     float z;
//     float w;
//     float l;
//     float h;
//     float rt;
//     // int id;
//     // float score;
//     Bndbox(){};
//     Bndbox(float x_, float y_, float z_, float w_, float l_, float h_, float rt_)
//         : x(x_), y(y_), z(z_), w(w_), l(l_), h(h_), rt(rt_) {}
// };


// class PostProcessCuda {
//   private:
//     Params params_;
//     float *anchors_;
//     float *anchor_bottom_heights_;
//     int *object_counter_;
//     cudaStream_t stream_ = 0;
//   public:
//     PostProcessCuda(cudaStream_t stream = 0);
//     ~PostProcessCuda();

//     int doPostprocessCuda(const float *cls_input, float *box_input, const float *dir_cls_input, float *bndbox_output);
// };

int nms_cpu(std::vector<Bndbox> bndboxes, const float nms_thresh, std::vector<Bndbox> &nms_pred);

namespace dsgn{
  struct PostProcessParams{
    static const int num_classes = 7;
    const char *class_name [num_classes] = {"car","truck_head","RTG","tray_w_container","QC","lock_station","street_lamp"};
    float min_x_range = 2.0;
    float max_x_range = 59.6;
    float min_y_range = -30.4;
    float max_y_range = 30.4;
    float min_z_range = -3.0;
    float max_z_range = 1.0;
    // the size of a voxel
    float voxel_x_size = 0.4;
    float voxel_y_size = 0.4;
    float voxel_z_size = 0.4;
    float dir_offset = 0.78539;
    float dir_limit_offset = 0.0;
    // the num of direction classes(bins)
    int num_dir_bins = 2;
    // anchors decode by (x, y, z, dir)
    int num_anchors = num_classes * 2;
    int len_per_anchor = 4;
    float anchors[56]{
      3.9, 1.6, 1.56, 0.0,
      3.9, 1.6, 1.56, 1.57,
      4.83, 3, 3.5, 0.0,
      4.83, 3, 3.5, 1.57,
      14, 1.89, 3, 0.0,
      14, 1.89, 3, 1.57,
      12.39, 2.74, 4, 0.0,
      12.39, 2.74, 4, 1.57,
      27.6, 3.87, 8, 0.0,
      27.6, 3.87, 8, 1.57,
      3.5, 2, 2.6, 0.0,
      3.5, 2, 2.6, 1.57,
      7.42, 2.81, 2.83, 0.0,
      7.42, 2.81, 2.83, 1.57
    }; 
    float anchor_bottom_heights[num_classes] = {-1.78,-1.78,-1.78, -1.78, -1.78, -1.78, -1.78};
    // the score threshold for classification
    float score_thresh = 0.1;
    float nms_thresh = 0.3; 
    // int max_num_pillars = MAX_VOXELS;
    // int pillarPoints_bev = max_num_points_per_pillar * max_num_pillars;
    // the detected boxes result decode by (x, y, z, w, l, h, yaw)
    int num_box_values = 7;
    // the input size of the 2D backbone network
    int grid_x_size = (max_x_range - min_x_range) / voxel_x_size; // 288
    int grid_y_size = (max_y_range - min_y_range) / voxel_y_size; // 304
    int grid_z_size = (max_z_range - min_z_range) / voxel_z_size; // 20
    // the output size of the 2D backbone network
    int feature_x_size = grid_x_size; // 288
    int feature_y_size = grid_y_size; // 304
    int nms_pre_maxsize = 4096;
    // PostProcessParams() {};
    };


  class PostProcessCuda {
    public:
      virtual std::vector<Bndbox> forward(float *cls_input, float *box_input, float *dir_cls_input, void* stream) = 0;
      virtual void print() = 0;

    private:
      std::vector<Bndbox> res_;
      std::vector<Bndbox> nms_pred_;
      std::vector<int> cls_ids_;
      std::vector<float> scores_;
    };
    
  
  std::shared_ptr<PostProcessCuda> create_postprocess(const PostProcessParams& params);
}

#endif
